import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 超参数设置
batch_size = 64
learning_rate = 0.001
epochs = 5

# 检查是否有可用的 GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

# 数据增强和标准化
train_transform = transforms.Compose([
    transforms.RandomRotation(10),         # 随机旋转±10度
    transforms.RandomAffine(0, translate=(0.1, 0.1)),  # 随机平移
    transforms.RandomResizedCrop(28, scale=(0.9, 1.1)), # 随机缩放后裁剪
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 测试数据只进行标准化
test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 加载 MNIST 数据集
train_dataset = datasets.MNIST(root='./data', train=True, transform=train_transform, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=test_transform, download=True)

train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

# 定义卷积神经网络模型
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 卷积层1：输入通道1，输出通道32，卷积核大小3x3
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        # 卷积层2：输入通道32，输出通道64，卷积核大小3x3
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        # 最大池化层
        self.pool = nn.MaxPool2d(2, 2)  # 池化窗口大小2x2
        # 全连接层1：输入为64*7*7，输出为128
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        # 全连接层2：输出为10（对应10类数字）
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # 第一层卷积 + ReLU + 池化
        x = self.pool(F.relu(self.conv1(x)))
        # 第二层卷积 + ReLU + 池化
        x = self.pool(F.relu(self.conv2(x)))
        # 展平特征图为一维向量
        x = x.view(-1, 64 * 7 * 7)
        # 全连接层1 + ReLU
        x = F.relu(self.fc1(x))
        # 全连接层2 (输出)
        x = self.fc2(x)
        return x

# 实例化模型并将其移到设备
model = CNN().to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 交叉熵损失用于分类问题
optimizer = optim.Adam(model.parameters(), lr=learning_rate)  # Adam 优化器

# 训练模型
def train(model, train_loader, criterion, optimizer, epochs, device):
    model.train()
    for epoch in range(epochs):
        running_loss = 0.0
        for batch_idx, (data, target) in enumerate(train_loader):
            # 将数据和目标移到设备（CPU/GPU）
            data, target = data.to(device), target.to(device)

            # 梯度清零
            optimizer.zero_grad()

            # 前向传播
            output = model(data)

            # 计算损失
            loss = criterion(output, target)

            # 反向传播
            loss.backward()

            # 更新参数
            optimizer.step()

            # 累计损失
            running_loss += loss.item()
            if batch_idx % 100 == 99:  # 每100个batch打印一次日志
                print(f'Epoch [{epoch+1}/{epochs}], Batch [{batch_idx+1}/{len(train_loader)}], Loss: {running_loss/100:.4f}')
                running_loss = 0.0

# 测试模型
def test(model, test_loader, device):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)  # 将数据移到设备
            output = model(data)
            _, predicted = torch.max(output.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item()

    print(f'Test Accuracy: {100 * correct / total:.2f}%')

# 训练与测试
train(model, train_loader, criterion, optimizer, epochs, device)
test(model, test_loader, device)

# 保存模型
torch.save(model.state_dict(), 'cnn_mnist_model.pth')
print("Model saved successfully.")
